import os  
import numpy as np  
import matplotlib.pyplot as plt  
from sko.PSO import PSO  
  
# 定义成本计算函数  
def calculate_total_cost(decision_vars, parameters):  
    x1, x2, xf, xd = decision_vars  
    p1, p2, C_assemble, C_finished, C_dismantle, C_exchange, Ct1, Ct2, n1, n2, n_finished = parameters  
      
    # 零配件检测成本  
    cost_parts = x1 * Ct1 * n1 + x2 * Ct2 * n2  
      
    # 装配成本  
    cost_assemble = C_assemble * n_finished  
      
    # 产品检测成本  
    cost_product = xf * C_finished * n_finished  
      
    # 成品次品率计算  
    p_assembly = (1 - x1) * p1 + (1 - x2) * p2 - (1 - x1) * (1 - x2) * p1 * p2  
    pf = 1 - (1 - p_assembly) ** 2  
      
    # 拆解和调换成本  
    cost_rework = xd * C_dismantle * pf * n_finished  
    cost_replace = (1 - xd) * C_exchange * pf * n_finished  
      
    # 总成本  
    total_cost = cost_parts + cost_assemble + cost_product + cost_rework + cost_replace  
      
    return total_cost  
  
# 粒子群优化算法  
def optimize_cost(parameters, count):  
    # PSO参数设置  
    n_variables = 4  # x1, x2, xf, xd  
    pso = PSO(func=lambda x: calculate_total_cost(x, parameters), dim=n_variables, pop=60, max_iter=50, lb=[0, 0, 0, 0], ub=[1, 1, 1, 1])  
      
    # 运行PSO  
    pso.run()  
      
    # 获取最优解  
    best_individual = [int(i) for i in pso.gbest_x]  
    best_cost = pso.gbest_y  
      
    print(f"Best Individual: {best_individual} (x1, x2, xf, xd)")  
    print(f"Minimum Cost: {best_cost}")  
      
    # 绘制迭代图  
    plt.figure(figsize=(10, 6))  
    plt.plot(pso.gbest_y_hist, marker='o', linestyle='-', color='b')  
    plt.title('Best Cost per Iteration')  
    plt.xlabel('Iteration')  
    plt.ylabel('Best Cost')  
    plt.grid(True)  
      
    # 确保输出文件夹存在  
    output_dir = 'test2_pso_picture'  
    if not os.path.exists(output_dir):  
        os.makedirs(output_dir)  
      
    # 保存图片到指定文件夹  
    plt.savefig(os.path.join(output_dir, f'best_cost_per_iteration{count}.png'))  
    plt.close()  
  
# 参数设置  
parameters_list = [  
    (0.1, 0.1, 6, 3, 5, 6, 2, 3, 100, 100, 80),  
    (0.2, 0.2, 6, 3, 5, 6, 2, 3, 100, 100, 80),  
    (0.1, 0.1, 6, 3, 5, 30, 2, 3, 100, 100, 80),  
    (0.2, 0.2, 6, 2, 5, 30, 1, 1, 100, 100, 80),  
    (0.1, 0.2, 6, 2, 5, 10, 8, 1, 100, 100, 80),  
    (0.05, 0.05, 6, 3, 40, 10, 2, 3, 100, 100, 80)  
]  
  
# 调用优化函数  
for index, parameters in enumerate(parameters_list, start=1):  
    optimize_cost(parameters, index)